import pygame

from src.base_elements.landscape import Home, LANDSCAPE_MAP, Steel, River


class MapGenerator:

    def __init__(self, stage_num):
        self.stage_num = stage_num
        self.map_file = f"assets/stages/stage_{stage_num}/map"
        self.map_data = self._load_map_data()
        self.landscape_group = pygame.sprite.Group()
        self.prop_forbidden_landscape_group = pygame.sprite.Group()

    def _load_map_data(self):
        with open(self.map_file) as f:
            lines = f.readlines()
            map_data = []
            for line in lines:
                map_data.append(line.split())
        return map_data

    def run(self, screen):
        for y, line in enumerate(self.map_data):
            for x, landscape_index in enumerate(line):
                if landscape_index == '0':
                    continue
                landscape_block = LANDSCAPE_MAP.get(landscape_index)()
                landscape_block.rect.left, landscape_block.rect.top = 3 + x * 24, 3 + y * 24
                self.landscape_group.add(landscape_block)
                if isinstance(landscape_block, Steel) or isinstance(landscape_block, River):
                    self.prop_forbidden_landscape_group.add(landscape_block)
        home_block = Home()
        self.landscape_group.add(home_block)
        self.prop_forbidden_landscape_group.add(home_block)
        self.landscape_group.draw(screen)
        pygame.display.flip()
        return self.landscape_group, self.prop_forbidden_landscape_group, home_block
